/*
 * Floppy Driver
 * main.c
 *
 * Copyright (C) 2019 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "floppy.h"
#include <memory.h>
#include <isa_dma.h>
#include <irq.h>
#include <thread.h>
#include <heap.h>
#include <semaphore.h>

const char driver_name[] = "floppy";

static dword_t floppy_motor_status[FLOPPY_DRIVES];
static DECLARE_LOCK(floppy_lock);
static DECLARE_LOCK(floppy_motor_lock);
static semaphore_t irq_mutex;
static thread_t *floppy_thread = NULL;

static dword_t floppy_init(void);
static dword_t floppy_cleanup(void);
static dword_t floppy_read(device_t *device, void *buffer, qword_t offset, size_t length, size_t *bytes_read);
static dword_t floppy_write(device_t *device, const void *buffer, qword_t offset, size_t length, size_t *bytes_written);
static dword_t floppy_ioctl(
    device_t *device,
    dword_t control_code,
    const void *in_buffer,
    size_t in_length,
    void *out_buffer,
    size_t out_length
);

static block_dev_driver_t floppy_driver_block =
{
    .init_proc = floppy_init,
    .cleanup_proc = floppy_cleanup,
    .read_proc = floppy_read,
    .write_proc = floppy_write,
    .ioctl_proc = floppy_ioctl
};

static void floppy_irq_handler(registers_t *regs, byte_t irq_num)
{
    UNUSED_PARAMETER(regs);
    UNUSED_PARAMETER(irq_num);

    release_mutex(&irq_mutex);
}

static dword_t floppy_wait_irq(dword_t timeout)
{
    return wait_mutex(&irq_mutex, timeout);
}

static bool_t floppy_fifo_write(byte_t byte)
{
    dword_t timeout = FLOPPY_IO_TIMEOUT;

    while (timeout--)
    {
        byte_t msr = cpu_read_port_byte(FLOPPY_MSR);
        if (msr & 0x80) break;
        syscall_sleep(1);
    }

    if (timeout != 0) cpu_write_port_byte(FLOPPY_FIFO, byte);
    else return FALSE;

    return TRUE;
}

static bool_t floppy_fifo_read(byte_t *byte)
{
    dword_t timeout = FLOPPY_IO_TIMEOUT;

    while (timeout--)
    {
        byte_t msr = cpu_read_port_byte(FLOPPY_MSR);
        if (msr & 0x80) break;
        syscall_sleep(1);
    }

    if (timeout != 0) *byte = cpu_read_port_byte(FLOPPY_FIFO);
    else return FALSE;

    return TRUE;
}

static word_t floppy_sense_interrupt(void)
{
    byte_t st0, cylinder;

    if (!floppy_fifo_write(FLOPPY_CMD_SENSE_INT)) return 0;
    if (!floppy_fifo_read(&st0)) return 0;
    if (!floppy_fifo_read(&cylinder)) return 0;

    return (st0 << 8) | cylinder;
}

static void floppy_select(byte_t drive)
{
    byte_t dor = cpu_read_port_byte(FLOPPY_DOR);
    dor &= 0xFC;
    dor |= drive & 1;
    cpu_write_port_byte(FLOPPY_DOR, dor);
}

static void floppy_reset()
{
    dword_t i;

    cpu_write_port_byte(FLOPPY_DOR, 0x00);
    cpu_write_port_byte(FLOPPY_DOR, 0x0C);
    floppy_wait_irq(NO_TIMEOUT);

    for (i = 0; i < 4; i++) floppy_sense_interrupt();

    floppy_fifo_write(FLOPPY_CMD_CONFIGURE);
    floppy_fifo_write(0);
    floppy_fifo_write(0x57);
    floppy_fifo_write(0);

    floppy_fifo_write(FLOPPY_CMD_SPECIFY);
    floppy_fifo_write(FLOPPY_SRT << 4 | FLOPPY_HUT);
    floppy_fifo_write(FLOPPY_HLT << 1);

    cpu_write_port_byte(FLOPPY_CCR, 0x00);
    floppy_select(0);
}

static dword_t floppy_motor_thread(void *param)
{
    dword_t i;

    while (TRUE)
    {
        lock_acquire(&floppy_motor_lock);

        for (i = 0; i < FLOPPY_DRIVES; i++)
        {
            if (floppy_motor_status[i] == FLOPPY_MOTOR_IDLE)
            {
                floppy_motor_status[i] = FLOPPY_MOTOR_STOPPING;
            }
            else if (floppy_motor_status[i] == FLOPPY_MOTOR_STOPPING)
            {
                byte_t dor = cpu_read_port_byte(FLOPPY_DOR);
                dor &= ~(1 << (4 + i));
                cpu_write_port_byte(FLOPPY_DOR, dor);

                floppy_motor_status[i] = FLOPPY_MOTOR_STOPPED;
            }
        }

        lock_release(&floppy_motor_lock);
        syscall_sleep(3000);
    }

    return 0;
}

static void floppy_motor_on(dword_t drive)
{
    lock_acquire(&floppy_motor_lock);

    if (floppy_motor_status[drive & 1] == FLOPPY_MOTOR_STOPPED)
    {
        byte_t dor = cpu_read_port_byte(FLOPPY_DOR);
        dor |= 1 << (4 + (drive & 1));
        cpu_write_port_byte(FLOPPY_DOR, dor);

        syscall_sleep(300);
    }

    floppy_motor_status[drive & 1] = FLOPPY_MOTOR_RUNNING;
    lock_release(&floppy_motor_lock);
}

static void floppy_motor_off(dword_t drive)
{
    lock_acquire(&floppy_motor_lock);
    floppy_motor_status[drive & 1] = FLOPPY_MOTOR_IDLE;
    lock_release(&floppy_motor_lock);
}

static bool_t floppy_recalibrate(byte_t drive)
{
    bool_t success = FALSE;
    int retries = 5;

    drive &= 1;

    floppy_select(drive);
    floppy_motor_on(drive);

    while (!success && ((retries--) > 0))
    {
        floppy_fifo_write(FLOPPY_CMD_RECALIBRATE);
        floppy_fifo_write(drive);
        floppy_wait_irq(FLOPPY_SEEK_TIMEOUT);

        byte_t st0 = floppy_sense_interrupt() >> 8;
        if (st0 & FLOPPY_ST0_SEEK_END) success = TRUE;
    }

    floppy_motor_off(drive);
    return success;
}

static dword_t floppy_init(void)
{
    init_mutex(&irq_mutex, 0);

    cpu_write_port_byte(CMOS_CMD_PORT, CMOS_FLOPPY_REG);
    byte_t floppy_data = cpu_read_port_byte(CMOS_DATA_PORT);
    if (!floppy_data) return ERR_NOTFOUND;

    dword_t ret = register_irq_handler(FLOPPY_IRQ, floppy_irq_handler, TRUE);
    if (ret != ERR_SUCCESS) goto cleanup;

    ret = create_system_thread(floppy_motor_thread, 0, THREAD_PRIORITY_LOW, 0, NULL, &floppy_thread);
    if (ret != ERR_SUCCESS) goto cleanup;

    floppy_reset();

    if ((floppy_data >> 4) == FLOPPY_TYPE_144MB)
    {
        floppy_recalibrate(0);

        device_t *device = (device_t*)malloc(sizeof(device_t));
        if (device == NULL)
        {
            ret = ERR_NOMEMORY;
            goto cleanup;
        }

        device->driver = &floppy_driver_block;
        strcpy(device->name, FLOPPY_FIRST_NAME);
        lock_init(&device->lock);
        device->capacity = FLOPPY_SECTOR_SIZE * FLOPPY_SECTORS;

        ret = register_block_device(device);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

    if ((floppy_data & 0x0F) == FLOPPY_TYPE_144MB)
    {
        floppy_recalibrate(1);

        device_t *device = (device_t*)malloc(sizeof(device_t));
        if (device == NULL)
        {
            ret = ERR_NOMEMORY;
            goto cleanup;
        }

        device->driver = &floppy_driver_block;
        strcpy(device->name, FLOPPY_SECOND_NAME);
        lock_init(&device->lock);
        device->capacity = FLOPPY_SECTOR_SIZE * FLOPPY_SECTORS;

        ret = register_block_device(device);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

cleanup:
    if (ret != ERR_SUCCESS)
    {
        if (floppy_thread) terminate_thread_internal(floppy_thread, 0);

        device_t *first = get_block_device(FLOPPY_FIRST_NAME);
        device_t *second = get_block_device(FLOPPY_SECOND_NAME);

        if (first)
        {
            unregister_block_device(first);
            free(first);
        }

        if (second)
        {
            unregister_block_device(second);
            free(second);
        }

        unregister_irq_handler(FLOPPY_IRQ, floppy_irq_handler);
    }

    return ret;
}

static dword_t floppy_cleanup(void)
{
    device_t *first = get_block_device(FLOPPY_FIRST_NAME);
    device_t *second = get_block_device(FLOPPY_SECOND_NAME);

    if (first)
    {
        unregister_block_device(first);
        free(first);
    }

    if (second)
    {
        unregister_block_device(second);
        free(second);
    }

    unregister_irq_handler(FLOPPY_IRQ, floppy_irq_handler);
    return floppy_thread ? terminate_thread_internal(floppy_thread, 0) : ERR_SUCCESS;
}

static dword_t floppy_read(device_t *device, void *buffer, qword_t offset, size_t length, size_t *bytes_read)
{
    dword_t ret = ERR_SUCCESS;
    dword_t count = 0;
    byte_t drive;
    byte_t cyl, head;
    dword_t sector_offset = (dword_t)offset % FLOPPY_SECTOR_SIZE;
    floppy_chs_t start_block = FLOPPY_LBA_TO_CHS((dword_t)offset / FLOPPY_SECTOR_SIZE);
    floppy_chs_t end_block = FLOPPY_LBA_TO_CHS((dword_t)(offset + length - 1) / FLOPPY_SECTOR_SIZE);
    byte_t *dma_buffer = NULL;

    if (offset >= (qword_t)(FLOPPY_SECTORS * FLOPPY_SECTOR_SIZE)) return ERR_INVALID;

    if (strcmp(device->name, FLOPPY_FIRST_NAME) == 0) drive = 0;
    else if (strcmp(device->name, FLOPPY_SECOND_NAME) == 0) drive = 1;
    else return ERR_INVALID;

    size_t max_transfer = ((dword_t)(offset + length - 1) / FLOPPY_SECTOR_SIZE) - ((dword_t)offset / FLOPPY_SECTOR_SIZE) + 1;
    if (max_transfer > FLOPPY_SPT) max_transfer = FLOPPY_SPT;

    area_t dma_area;
    ret = isa_dma_alloc(max_transfer * FLOPPY_SECTOR_SIZE, &dma_area);
    if (ret != ERR_SUCCESS) return ret;

    ret = memory_view_area(memory_upper_space, (void**)&dma_buffer, &dma_area, MEMORY_FLAG_ACCESSIBLE);
    if (ret != ERR_SUCCESS)
    {
        isa_dma_free(&dma_area);
        return ret;
    }

    lock_acquire(&floppy_lock);
    floppy_select(drive);
    floppy_motor_on(drive);

    for (cyl = start_block.cylinder; cyl <= end_block.cylinder; cyl++)
    {
        byte_t first_head = 0, last_head = 1;

        if (cyl == start_block.cylinder) first_head = start_block.head;
        if (cyl == end_block.cylinder) last_head = end_block.head;

        for (head = first_head; head <= last_head; head++)
        {
            bool_t success = FALSE;
            byte_t retries = 3;
            byte_t first_sector = 1;
            byte_t last_sector = FLOPPY_SPT;

            if ((cyl == start_block.cylinder) && (head == first_head)) first_sector = start_block.sector;
            if ((cyl == end_block.cylinder) && (head == last_head)) last_sector = end_block.sector;
            dword_t amount = (last_sector - first_sector + 1) * FLOPPY_SECTOR_SIZE;

            while (retries--)
            {
                byte_t i;
                byte_t command_input[FLOPPY_RW_INPUTS] =
                {
                    FLOPPY_CMD_READ | FLOPPY_FLAG_MFM,
                    drive | (head << 2),
                    cyl,
                    head,
                    first_sector,
                    0x02,
                    last_sector,
                    0x1B,
                    0xFF
                };
                byte_t command_output[FLOPPY_RW_OUTPUTS];

                isa_dma_write(FLOPPY_DMA_CHANNEL, &dma_area, amount);

                for (i = 0; i < FLOPPY_RW_INPUTS; i++) if (!floppy_fifo_write(command_input[i])) goto timeout;
                if (floppy_wait_irq(FLOPPY_READ_TIMEOUT) == ERR_TIMEOUT) goto timeout;
                for (i = 0; i < FLOPPY_RW_OUTPUTS; i++) if (!floppy_fifo_read(&command_output[i])) goto timeout;

                if ((command_output[0] & 0xC0) || command_output[1] || command_output[2]) continue;

                success = TRUE;
                break;

            timeout:
                floppy_reset();
                floppy_recalibrate(drive);
            }

            if (!success)
            {
                ret = ERR_HARDWARE;
                goto done;
            }

            memcpy(&((byte_t*)buffer)[count], &dma_buffer[sector_offset], MIN(length - count, amount - sector_offset));
            count += MIN(length - count, amount - sector_offset);
            sector_offset = 0;
        }
    }

done:
    floppy_motor_off(drive);
    lock_release(&floppy_lock);
    memory_free(memory_upper_space, dma_buffer);
    isa_dma_free(&dma_area);

    if (bytes_read) *bytes_read = count;
    return ret;
}

static dword_t floppy_write(device_t *device, const void *buffer, qword_t offset, size_t length, size_t *bytes_written)
{
    dword_t ret = ERR_SUCCESS;
    dword_t count = 0;
    byte_t drive;
    byte_t cyl, head;
    dword_t sector_offset = (dword_t)(offset % FLOPPY_SECTOR_SIZE);
    floppy_chs_t start_block = FLOPPY_LBA_TO_CHS((dword_t)offset / FLOPPY_SECTOR_SIZE);
    floppy_chs_t end_block = FLOPPY_LBA_TO_CHS((dword_t)(offset + length - 1) / FLOPPY_SECTOR_SIZE);
    byte_t first_sector_data[FLOPPY_SECTOR_SIZE];
    byte_t last_sector_data[FLOPPY_SECTOR_SIZE];
    byte_t *dma_buffer = NULL;

    if (offset >= (qword_t)(FLOPPY_SECTORS * FLOPPY_SECTOR_SIZE)) return ERR_INVALID;

    if (strcmp(device->name, FLOPPY_FIRST_NAME) == 0) drive = 0;
    else if (strcmp(device->name, FLOPPY_SECOND_NAME) == 0) drive = 1;
    else return ERR_INVALID;

    if (sector_offset)
    {
        ret = floppy_read(device, first_sector_data, offset & ~(FLOPPY_SECTOR_SIZE - 1), FLOPPY_SECTOR_SIZE, NULL);
        if (ret != ERR_SUCCESS) return ret;
    }

    if ((offset + length) % FLOPPY_SECTOR_SIZE)
    {
        ret = floppy_read(device, last_sector_data, (offset + length) & ~(FLOPPY_SECTOR_SIZE - 1), FLOPPY_SECTOR_SIZE, NULL);
        if (ret != ERR_SUCCESS) return ret;
    }

    size_t max_transfer = ((dword_t)(offset + length - 1) / FLOPPY_SECTOR_SIZE) - ((dword_t)offset / FLOPPY_SECTOR_SIZE) + 1;
    if (max_transfer > FLOPPY_SPT) max_transfer = FLOPPY_SPT;

    area_t dma_area;
    ret = isa_dma_alloc(max_transfer * FLOPPY_SECTOR_SIZE, &dma_area);
    if (ret != ERR_SUCCESS) return ret;

    ret = memory_view_area(memory_upper_space, (void**)&dma_buffer, &dma_area, MEMORY_FLAG_ACCESSIBLE | MEMORY_FLAG_WRITABLE);
    if (ret != ERR_SUCCESS)
    {
        isa_dma_free(&dma_area);
        return ret;
    }

    lock_acquire(&floppy_lock);
    floppy_select(drive);
    floppy_motor_on(drive);

    for (cyl = start_block.cylinder; cyl <= end_block.cylinder; cyl++)
    {
        byte_t first_head = 0, last_head = 1;

        if (cyl == start_block.cylinder) first_head = start_block.head;
        if (cyl == end_block.cylinder) last_head = end_block.head;

        for (head = first_head; head <= last_head; head++)
        {
            bool_t success = FALSE;
            byte_t retries = 3;
            byte_t first_sector = 1;
            byte_t last_sector = FLOPPY_SPT;

            if ((cyl == start_block.cylinder) && (head == first_head)) first_sector = start_block.sector;
            if ((cyl == end_block.cylinder) && (head == last_head)) last_sector = end_block.sector;
            dword_t amount = (last_sector - first_sector + 1) * FLOPPY_SECTOR_SIZE;

            if (sector_offset)
            {
                memcpy(dma_buffer, first_sector_data, FLOPPY_SECTOR_SIZE);
            }

            if ((length - count) < amount)
            {
                memcpy(&dma_buffer[amount - FLOPPY_SECTOR_SIZE], last_sector_data, FLOPPY_SECTOR_SIZE);
            }

            memcpy(&dma_buffer[sector_offset], &((const byte_t*)buffer)[count], MIN(length - count, amount - sector_offset));

            while (retries--)
            {
                byte_t i;
                byte_t command_input[FLOPPY_RW_INPUTS] =
                {
                    FLOPPY_CMD_WRITE | FLOPPY_FLAG_MFM,
                    drive | (head << 2),
                    cyl,
                    head,
                    first_sector,
                    0x02,
                    last_sector,
                    0x1B,
                    0xFF
                };
                byte_t command_output[FLOPPY_RW_OUTPUTS];

                isa_dma_read(FLOPPY_DMA_CHANNEL, &dma_area, amount);

                for (i = 0; i < FLOPPY_RW_INPUTS; i++) if (!floppy_fifo_write(command_input[i])) goto write_timeout;
                if (floppy_wait_irq(FLOPPY_READ_TIMEOUT) == ERR_TIMEOUT) goto write_timeout;
                for (i = 0; i < FLOPPY_RW_OUTPUTS; i++) if (!floppy_fifo_read(&command_output[i])) goto write_timeout;

                if ((command_output[0] & 0xC0) || command_output[1] || command_output[2]) continue;

                success = TRUE;
                break;

            write_timeout:
                floppy_reset();
                floppy_recalibrate(drive);
            }

            if (!success)
            {
                ret = ERR_HARDWARE;
                goto done;
            }

            count += MIN(length - count, amount - sector_offset);
            sector_offset = 0;
        }
    }

done:
    floppy_motor_off(drive);
    lock_release(&floppy_lock);
    memory_free(memory_upper_space, dma_buffer);
    isa_dma_free(&dma_area);

    if (bytes_written) *bytes_written = count;
    return ret;
}

static dword_t floppy_ioctl(device_t *device,
                            dword_t control_code,
                            const void *in_buffer,
                            size_t in_length,
                            void *out_buffer,
                            size_t out_length)
{
    switch (control_code)
    {
    case IOCTL_BLOCK_DEV_MEDIA_INFO:
        if (out_length >= sizeof(block_dev_media_info_t))
        {
            ((block_dev_media_info_t*)out_buffer)->heads = 2;
            ((block_dev_media_info_t*)out_buffer)->tracks = 80;
            ((block_dev_media_info_t*)out_buffer)->sectors_per_track = 18;
            ((block_dev_media_info_t*)out_buffer)->bytes_per_sector = 512;

            return ERR_SUCCESS;
        }
        else
        {
            return ERR_SMALLBUF;
        }

        break;

    default:
        return ERR_INVALID;
    }
}

dword_t driver_load(const char *parameters)
{
    return register_block_dev_driver(&floppy_driver_block);
}
